# given contour and reference output a image
import os
from argparse import ArgumentParser

from torch.utils.data import Dataset
import torch
import torchvision.transforms as T
import torchvision.utils as vutils
from PIL import Image

from models.networks.model import Model
from option import Options


class EvaluationDataset(Dataset):
    def __init__(self, root):
        super(EvaluationDataset, self).__init__()
        self.root = root
        self.contour_paths, self.image_paths = self.get_paths()
        self.contour_paths.sort()
        self.image_paths.sort()

    def __len__(self):
        return len(min(self.contour_paths,self.image_paths))

    def __getitem__(self, idx):
        sketch = self.contour_paths[idx]
        ref = self.image_paths[idx]
        sketch = Image.open(sketch).convert('L')
        ref = Image.open(ref).convert('RGB')
        sketch = self.translation(sketch, True)
        ref = self.translation(ref, False)
        return sketch, ref

    def get_paths(self):
        contour_dir = os.path.join(self.root, 'contour')
        img_dir = os.path.join(self.root, 'ref')
        contour_paths = [os.path.join(contour_dir, f) for f in os.listdir(contour_dir)]
        image_paths = [os.path.join(img_dir, f) for f in os.listdir(img_dir)]
        return contour_paths, image_paths

    def translation(self, img, isSketch=True):
        if isSketch:
            t = T.Compose([T.Resize(256), T.ToTensor(), T.Normalize(0.5, 0.5)])
        else:
            t = T.Compose([T.Resize(256), T.ToTensor(), T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
        return t(img)

class ModelUtil:
    def __init__(self, opt, dataset):
        super(ModelUtil, self).__init__()
        self.opt = opt
        self.loader = self.getloader(dataset)
        self.save_root = os.path.join(os.path.dirname(opt.checkpoints_dir), 'output', opt.name)
        if not os.path.exists(self.save_root):
            os.makedirs(self.save_root)
        self.device = opt.device
        model = Model(opt)
        self.model = model

    def eval(self):
        self.model.eval()
        for i, data in enumerate(self.loader):
            sketch = data[0].to(self.opt.device)
            ref = data[1].to(self.opt.device)
            fake = self.model.inference(sketch, ref)['fake_image']
            try:
                vutils.save_image(fake, self.save_root + '/' + str(i) + '_' + 'fake' + '.png',
                                  nrow=1, padding=0, normalize=True)
                vutils.save_image(sketch, self.save_root + '/' + str(i) + '_' + 'sketch' + '.png',
                                  nrow=1, padding=0, normalize=True)
                vutils.save_image(ref, self.save_root + '/' + str(i) + '_' + 'ref' + '.png',
                                  nrow=1, padding=0, normalize=True)
            except OSError as err:
                print(err)


    def getloader(self, dataset):
        loader = torch.utils.data.DataLoader(dataset, batch_size=self.opt.batch_size,
                                             shuffle=False, drop_last=False, num_workers=self.opt.num_workers)
        return loader

if __name__ == '__main__':
    # parser = ArgumentParser()
    # parser.add_argument('--dataset',type=str, default=r'D:\zhw\data\TestDataset', help='test data root')
    # parser.add_argument('--batch_size', type=int, default=1)
    # parser.add_argument('--num_workers', type=int, default=0)
    # parser.add_argument('--checkpoints_pth', type=str, default='./latest_G_old.pth', help='path of ckp')
    # parser.add_argument('--name',type=str,default='val', help='name of experiment')
    # parser.add_argument('--device', type=str, default='cpu')
    # args = parser.parse_args()
    # print(vars(args))
    train_parser = Options()
    opt = train_parser.parse()
    dataset = EvaluationDataset(root=opt.dataroot)
    model_util = ModelUtil(opt, dataset)
    model_util.eval()
